import logging
import logging.config
import HomologFinderUtility, FileUtility
import random
import os
import re
import glob
import sys
from STRINGUtility import STRINGUtility,filterSequencesByCOG
from Bio.Blast import NCBIStandalone
from Bio import SeqIO
from datamodel import HIK,HmmerHit,HmmerHitValidation, HomologFinderConfig, HomologSearch,BlastHit,BlastHit,Homolog,Homolog, DatabaseObj, NoSequencesToExtractError
import threading
from Queue import Queue
from optparse import OptionParser
logging.config.fileConfig("logging.conf")
logger = logging.getLogger("MAIN_LOG")

def prepareHIK(hik):
    '''Makes sure all parts are in place for the homolog identity kit to be used on the run'''
    hik.hmmFile = HomologFinderConfig.HMM_STORE + hik.hmm + ".hmm"
    hik.stringHmm = HomologFinderConfig.HMM_STORE + hik.cog + ".string.hmm"
    hik.stringFasta = HomologFinderConfig.SCRATCH + hik.cog + ".fasta"
    hik.consensus = HomologFinderConfig.CONSENSUS_STORE + hik.cog + ".cs"
    hik.proteinProfile = HomologFinderConfig.PP_STORE + hik.cog + ".chk"

    if (not FileUtility.safeFileExists(hik.proteinProfile)):
        try:
            createHmmFromSTRINGDB(hik)
            HomologFinderUtility.makeConsensusFileFromHmm(hik.stringHmm, hik.consensus)
            HomologFinderUtility.createBlastDb(hik.stringFasta, hik.stringFasta)
            makeProteinProfile(hik, processors=2)
        except Exception as e:
            raise e


def createHmmFromSTRINGDB(hik):
    """Creates an hmm from the String database from sequences which:
        1. match the Pfam hmm assigned to the gene
        2. share the COG with the homolog
    """
    if (not FileUtility.safeFileExists(hik.stringHmm)):
        try:
            logger.debug("Creating an Hmm for %s from the STRING db - outputting to %s" % (hik.cog, hik.stringHmm))
            makeFastaFileFromString(hik)
            alnFile = HomologFinderConfig.SCRATCH + hik.cog + ".aln"
            alignFastasUsingMafft(hik.stringFasta, alnFile)
            stockAlnFile = HomologFinderConfig.SCRATCH + hik.cog + ".stock.aln"
            HomologFinderUtility.makeStockholmAlignment(alnFile, stockAlnFile)
            HomologFinderUtility.makeHmm(stockAlnFile, hik.stringHmm)
            FileUtility.safeFileRemove([stockAlnFile, alnFile])
        except NoSequencesToExtractError as e:
            raise e


def alignFastasUsingMafft(fileIn, fileOut):
    """Aligns sequences in a fasta using the mafft program"""
    if(not FileUtility.safeFileExists(fileOut)):
        logger.info("Aligning sequences using mafft to %s" % fileOut)
        str = HomologFinderConfig.MAFFT + ' --thread %i --genafpair --maxiterate 1000 --quiet %s > %s' % (int(HomologFinderConfig.MAFFT_THREAD), fileIn, fileOut)
        os.system(str)


def makeFastaFileFromString(hik):
    if(not FileUtility.safeFileExists(hik.stringFasta)):
        logger.debug("Pulling fastas out of STRING for hmm creation")
        StringUtility = STRINGUtility()
        seqsMatchingCOG = StringUtility.getSequencesMatchingCOG(hik.cog)
        logger.info("Found %i sequences matching %s" % (len(seqsMatchingCOG), hik.cog))
        seqsMatchingCOG = set(seqsMatchingCOG)
        logger.info("Of which, %i were unique" % len(seqsMatchingCOG))

        #Now need to remove ones which have multiple cogs associated with them
        multiCOG = StringUtility.getMultiCOGSequences(seqsMatchingCOG)
        singleCOG = [p for p in seqsMatchingCOG if p not in multiCOG]

        #fastas = StringUtility.getFastaList(singleCOG, HomologFinderConfig.STRING_SEQS)
        tmpFile = HomologFinderConfig.SCRATCH + '%016x.tmp' % random.getrandbits(128)
        try:
            fastas = HomologFinderUtility.getSequencesFromDb(singleCOG, STRINGUtility.STRING_DB, tmpFile)
            logger.info("Checking the selected sequences with a signature hmm (%s) using hmmsearch" % hik.hmm)
            HomologFinderUtility.getHmmFromPfam(hik.hmm, hik.hmmFile)
            tmpFile2 = HomologFinderConfig.SCRATCH + '%016x.tmp' % random.getrandbits(128)
            HomologFinderUtility.runHmmer(tmpFile, hik.hmmFile, tmpFile2)
            hmmerHits = HomologFinderUtility.parseHmmerOutput(tmpFile2)
            logger.info("Have found %i homologs in remaining sequences for %s" % (len(hmmerHits[HmmerHitValidation.HOMOLOGOUS]), hik.cog))
            seqs = list()
            if not hmmerHits:
                raise NoSequencesToExtractError('Hmmer sequence search on COG homologs for %s returned no hits' % hik.cog)
            for hit in hmmerHits[HmmerHitValidation.HOMOLOGOUS]:
                seqs.append(hit.getSequence())

            if (len(seqs) > 500):
                logging.debug("Over a 500 sequences - randomly subsampling")
                seqs = random.sample(seqs, 500)
            fastas = StringUtility.getFastaList(seqs, tmpFile)
            logger.info("Creating fasta file from STRING for %s consisting of %i sequences at %s" %(hik.cog, len(seqs), hik.stringFasta))
            HomologFinderUtility.writeFastaFile(fastas, hik.stringFasta)

        except NoSequencesToExtractError as e:
           raise e
        finally:
            FileUtility.safeFileRemove([tmpFile, tmpFile2])




def readGeneList(geneList):
    """Create a list of hiks from the gene list"""

    f = open(geneList, 'r',1)
    hiks = list()
    for line in f:
        match = re.match('([^#]\S+)\t(\S+)\t(\S+)', line)
        if match:
            hik = HIK()
            hik.geneName = match.group(1)
            hik.cog = match.group(2)
            hik.hmm = match.group(3)
            hiks.append(hik)
    f.close()
    return hiks

def readDbList(dbList):
    """Create a list of dbs from the db list"""

    f = open(dbList, 'r',1)
    dbList = list()
    for line in f:
        match = re.match('([^#]\S+)\t(\S+)', line)
        if match:
            db = DatabaseObj()
            db.name = match.group(1)
            db.location = match.group(2)
            dbList.append(db)
    f.close()
    return dbList


def makeProteinProfile(hik, processors=16):
    '''Makes a protein profile from a sequence using blastpgp'''
    if (not FileUtility.safeFileExists(hik.proteinProfile)):
        reportFile = hik.proteinProfile + ".report"
        str = HomologFinderConfig.BLAST_BLASTPGP + " -i %s -a %i -d %s -j 2 -F 'm S' -C %s > %s" % (hik.consensus, processors, hik.stringFasta, hik.proteinProfile, reportFile)
        os.system(str)
        FileUtility.safeFileRemove(glob.glob(hik.stringFasta +'.p*'))
    else:
        logger.debug("Protein Profile %s already exists" % hik.proteinProfile)



class HomologyFinderWorker(threading.Thread):
    '''This is a class which represents a homology search
    i.e. looking for homologs against a specific database.
    '''
    def __init__(self, queue, processors=16, blast=False):
        threading.Thread.__init__(self)
        self.queue = queue
        self.processors = processors
        self.blast = blast

    def run(self):
        while True:
            try:
                hs = self.queue.get()
                hik, db = hs.hik, hs.db
                prepareHIK(hik)

                if HomologFinderUtility.isSearchPickled(hs):
                    logger.debug('Run of %s vs %s is already pickled' % (hik.geneName, db.name))
                elif not self.blast:
                        logger.info("Not performing the blast - just preparing the files")
                else:
                    logger.info('Looking for homologs of %s in %s' % (hik.geneName, db.name))
                    psiBlastOut = HomologFinderConfig.SCRATCH + "%s.vs.%s.psitblastn.out" % (hik.geneName, db.name)
                    if (not FileUtility.safeFileExists(psiBlastOut)):
                        logger.info("Performing psitblastn of %s vs %s" % (hik.geneName, db.name))
                        str = HomologFinderConfig.BLAST_BLASTALL + " -p psitblastn -a %i -b %i -v %i -R %s -F 'm S' -i %s -d %s -o %s -z 1000000000" % (self.processors, HomologFinderConfig.MAX_HOMOLOG_COUNT, HomologFinderConfig.MAX_HOMOLOG_COUNT, hik.proteinProfile, hik.consensus, db.location, psiBlastOut)
                        os.system(str)
                    else:
                        logger.debug("%s already exists" % psiBlastOut)
                    homologFasta = HomologFinderConfig.SCRATCH + "%s.%s.homologs.fasta" % (db.name, hik.geneName)
                    if (not FileUtility.safeFileExists(homologFasta)):
                        homologs = HomologFinderUtility.getHomologListFromBLASTOutput(psiBlastOut)
                        if (homologs):
                            HomologFinderUtility.getSequencesFromDb(homologs, db.location, homologFasta)

                    elif (os.path.getsize(homologFasta)):
                        logger.debug("%s already exists" % homologFasta)
                    else:
                        logger.debug("%s already exists, but has zero size" %homologFasta)

                    #Now we're going to perform a reciprocal blast against the STRING database

                    reciprocalBlastOut = HomologFinderConfig.SCRATCH + "%s.%s.homologs.vs.STRING.tab" % (db.name,hik.geneName )
                    tmpFile = HomologFinderConfig.SCRATCH + '%016x.tmp' % random.getrandbits(128)
                    if (not FileUtility.safeFileExists(reciprocalBlastOut)):
                        logger.info("Performing reciprocal blastx of %s homologs from %s vs STRING - output going to %s" % (hik.geneName, db.name, reciprocalBlastOut))
                        str = HomologFinderConfig.BLAST_BLASTALL + " -p blastx -a %i -b 1 -v 1 -d %s -o %s -i %s -m 8" % (self.processors, STRINGUtility.STRING_DB, reciprocalBlastOut, homologFasta)
                        os.system(str)

                    else:
                        logger.debug("%s already exists" % psiBlastOut)

                    #We'll parse this data and check to see that the best hit in STRING shares the original COG grouping
                    #If it does, we need to pull out the species for some taxonomy stuff
                    handle = None
                    try:
                        logger.info("Parsing reciprocal blast file %s" % reciprocalBlastOut)
                        handle = open(reciprocalBlastOut, 'r',1)
                        checkedFrags = list()
                        okFrags = list()
                        for line in handle:
                            hit = BlastHit(line)
                            #May get two hits per frag, but first one is the highest e-value and thus the longest
                            if hit.query not in checkedFrags:
                                okFrags.append(hit)
                                checkedFrags.append(hit.query)
                        logger.info("Found %i unique sequences in the file" % len(checkedFrags))
                        fragList = set()
                        for hit in okFrags:
                            fragList.add(hit.subject)
                        logger.info("These sequences match up to %i proteins in STRING" % len(fragList))
                        goodMatches = filterSequencesByCOG(fragList, hik.cog)
                        logger.info("Of these, %i match the COG %s" % (len(goodMatches), hik.cog))
                        badMatches = [p for p in fragList if p not in goodMatches]
                        logger.info("%i did not match the COG %s" % (len(badMatches), hik.cog))
                        for hit in okFrags:
                            homolog = Homolog()
                            homolog.fragment = hit.query
                            homolog.stringHit = hit.subject
                            homolog.species = hit.species
                            if hit.subject in goodMatches:
                                hs.goodHits.append(homolog)
                            else:
                                hs.badHits.append(homolog)

                        #pickle the search
                        HomologFinderUtility.pickleSearch(hs)

                    finally:
                        FileUtility.safeFileClose(handle)
            except Exception as e:
                logger.exception('An error has occured in HomologFinderWorker - %s' % e.value)
            self.queue.task_done()



def main():
    parser = OptionParser()
    parser.add_option('-g', '--geneList', dest="geneList", help="The list of the genes to be used", metavar='FILE')
    parser.add_option('-d', '--dbList', dest='dbList', help='The list of databases to run genes against', metavar='FILE')
    (options, args) = parser.parse_args()
    hiks = readGeneList(options.geneList)
    dbs = readDbList(options.dbList)

    q = Queue()
    for i in range(int(HomologFinderConfig.WORKER_COUNT)):
        t = HomologyFinderWorker(q, processors=int(HomologFinderConfig.PROCESSORS), blast=HomologFinderConfig.PERFORM_BLAST)
        t.setDaemon(True)
        t.start()
    for db in dbs:
        for hik in hiks:
            hs = HomologSearch(hik, db)
            q.put(hs)
            pass
    q.join()

if __name__ == "__main__":
    main()
